package com.putprize.percy

import scala.math.{log,random,abs}
import scala.io.Source

import org.apache.log4j.Logger

import java.io.PrintWriter

trait Model {
  def K:Int
  def M:Int
  def pvzBeta(v:Int, z:Int):Double
  def logPvzBeta(v:Int, z:Int):Double
  def INIT_T_VALUE:Double
  def INIT_V_VALUE:Double
}

// For Inference, Only Used in Predict

class Model1(
    val countT: Array[Double],
    val countV: Array[Map[Int,Double]],
    val initT:Double,
    val initV:Double
    ) extends Model {
  
  def K = countT.size
  def M = countV.size
  
  def pvzBeta(v:Int,z:Int) = {
    (countV(v).getOrElse(z, 0.0)+INIT_V_VALUE)/(countT(z)+countV.size*INIT_V_VALUE)
  }
  
  def logPvzBeta(v:Int,z:Int) = {
    log(countV(v).getOrElse(z, 0.0)+INIT_V_VALUE)-log(countT(z)+countV.size*INIT_V_VALUE)
  }  
  
  def INIT_T_VALUE = initT
  def INIT_V_VALUE = initV
}

// For Estimate, Only Used in Estimate

class Model2(
    val countT:Array[Double],
    val countV:Array[Array[Double]],
    val initT:Double,
    val initV:Double
    ) extends Model {
  val _log = Logger.getLogger(this.getClass.getName)
  
  def K = countT.size
  def M = countV.size

  def pvzBeta(v:Int,z:Int) = {
    if (v < 0 || v >= countV.size){
      _log.error("v "+v)
      sys.exit(1)
    }
    if (z < 0 || z >= countT.size){
      _log.error("z "+z)
      sys.exit(1)
    }
    (countV(v)(z)+INIT_V_VALUE)/(countT(z)+countV.size*INIT_V_VALUE)
  }
  
  def logPvzBeta(v:Int,z:Int) = {
    if (v < 0 || v >= countV.size){
      _log.error("v "+v)
      sys.exit(1)
    }
    if (z < 0 || z >= countT.size){
      _log.error("z "+z)
      sys.exit(1)
    } 
    log(countV(v)(z)+INIT_V_VALUE) - log(countT(z)+countV.size*INIT_V_VALUE)
  }  
  
  def INIT_T_VALUE = initT
  def INIT_V_VALUE = initV

  def getCountT(z:Int) = countT(z)
  def getCountV(v:Int,z:Int) = countV(v)(z)
  
  def setCountT(z:Int,c:Double) {
    countT(z) = c
  } 
  def setCountV(v:Int,z:Int,c:Double) {
    countV(v)(z) = c
  }
  
  def sumCountV(v:Int) ={
    countV(v).reduce(_+_)
  }
  
  def sumCountV() = {
    def _sum(A1:Array[Double], A2:Array[Double]) = {
      (0 until K).map( k => A1(k)+A2(k)).toArray
    }
    countV.par.reduce(_sum)
  }
}

object PyotrModel {
  
  val INIT_T_VALUE = 1.0
  val INIT_V_VALUE = 1.0
  
  private
  val _log = Logger.getLogger(this.getClass.getName)
  
  def initModel(K:Int, M:Int) = {
    _log.info("Init CountT ...")
    val countT = new Array[Double](K)
    for (t <- 0 until K){
      countT(t) = 0.0
    }
    _log.info("Init CountT Done")
    _log.info("Init CountV ...")
    val countV = new Array[Array[Double]](M)
    for (v <- 0 until M){
      countV(v) = new Array[Double](K)
      for (z <- 0 until K){
        countV(v)(z) = 1.0/M+random
        countT(z) += countV(v)(z) // 确实还要归一化
      }
    }
    _log.info("Init CountV Done")
    
    _log.info("INIT_T_VALUE "+INIT_T_VALUE/K)
    _log.info("INIT_V_VALUE "+INIT_V_VALUE/K)
    new Model2(countT,countV,INIT_T_VALUE/K,INIT_V_VALUE/K)
  }
  
  def initModel(K:Int,M:Int,Xs:Map[String,Document]) = {
    _log.info("Init Count ...")
    val countT = new Array[Double](K)
    for (z <- 0 until K){
      countT(z) = 0.0
    }
    val countV = new Array[Array[Double]](M)
    for (v <-0 until M){
      countV(v) = new Array[Double](K)
      for (z <- 0 until K){
        countV(v)(z) = 0.0
      }
    }
    _log.info("Init Count Done")
    _log.info("Count ...")
    var cnt = 0
    var rcnt = 1000000
    //var rindex = -1
    var randomArray = new Array[Double](rcnt + 1)
    (0 until rcnt).foreach{i =>
      randomArray(i) = random
    }
    Xs.foreach(iv => {
      cnt += 1
      if (cnt % 10000 == 0)
          _log.info("Count "+cnt)
      val doc = iv._2
      (0 until doc.n).par.foreach{ i =>
        val v = doc.vs(i)
        val c = doc.cs(i)
        val rs = new Array[Double](K)
        (0 until K).foreach {j =>
//          if (rindex == rcnt)
//        	rindex = 0
//          else
//            rindex += 1  
          val rindex = abs((cnt*i+j)%rcnt)
          rs(j) = randomArray(rindex)
        }
        val srs = rs.reduce(_+_)
        for (j <- 0 until K)
          rs(j) = rs(j)/srs
        (0 until K).foreach(z => {
          val s = c*rs(z)
          countV(v)(z) += s
          countT(z) += s
        })
      }
    })
    _log.info("Count Done")
    
    _log.info("INIT_T_VALUE "+INIT_T_VALUE/K)
    _log.info("INIT_V_VALUE "+INIT_V_VALUE/K)
    
    new Model2(countT,countV,INIT_T_VALUE/K,INIT_V_VALUE/K)
  }
 
  def initModel(K:Int, M:Int, V:Map[Int,Double],location:String) = {
    
    val countT = new Array[Double](K)
    for (z <- 0 until K){
      countT(z) = 0.0
    }
    val countV = new Array[Array[Double]](M)
    for (v <- 0 until M){
      countV(v) = new Array[Double](K)
      for (z <- 0 until K){
        countV(v)(z) = 0.0
      }
    }
    
    val source = Source.fromFile(location,"UTF-8")
    val data = source.getLines.toArray

    (0 until data.size).foreach { i =>
      if (i % 10000 == 0){
        _log.info(i)
      }
      val ss = V.getOrElse(i, 0.0)
      val line = data(i)
      if (line.trim() != ""){
	      val zs = line.trim().split(" ")
	      val us = zs.map(s => s.split(":")).map(s => {(s(0).toInt,s(1).toDouble)}).toMap
	      val rs = new Array[Double](K)
	      (0 until K).foreach {j =>
          	rs(j) = us.getOrElse(j,0.0)
	      }	  
	      val srs = rs.reduce(_+_)
	      (0 until K).foreach { z => {
	        val s = ss*rs(z)/srs
	        countV(i)(z) += s
	        countT(z) += s 
	      }}      
      } 
      else {
	      val rs = new Array[Double](K)
	      (0 until K).foreach {j =>
          	rs(j) = random
	      }	 
	      val srs = rs.reduce(_+_)
	      (0 until K).foreach {z =>
	        val s = ss*rs(z)/srs
	        countV(i)(z) += s
	        countT(z) += s
	      }
      }
    }
    source.close
    
    (0 until K).foreach{z => 
      _log.info(z+" "+countT(z))
    }
    
    _log.info("INIT_T_VALUE "+INIT_T_VALUE/K)
    _log.info("INIT_V_VALUE "+INIT_V_VALUE/K)
    val m = new Model2(countT,countV,INIT_T_VALUE/K,INIT_V_VALUE/K)
    m
  }
 
  def initModel(countT:Array[Double],
                countV:Array[Array[Double]],
                S:Double,
                N:Int) = {
    
    //val s = Dirichlet.optimize(S, countT.size, N)
    new Model2(countT,countV, INIT_T_VALUE/countT.size, INIT_V_VALUE/countT.size)
  }
  
  
  def saveCountT(countT:Array[Double],location:String) {
    val out = new PrintWriter(location)
    for (i <- 0 until countT.size){
      out.write(countT(i)+"\n")
    }
    out.close
  }
  
  def saveCountV(countT:Array[Double],countV:Array[Array[Double]],location:String) {
    val out = new PrintWriter(location)
    for (i <- 0 until countV.size){
      val vs = (0 until countV(i).size).par.
      			map( z => {(z,countV(i)(z))}).
      			filter( zv => (zv._2/countT(zv._1)) > 1e-7)
      val line = 
        if (vs.size > 0) 
        	vs.map( v => {v._1.toString+":"+v._2.toString}).reduce(_+" "+_)
        else
        	""
      out.write(line+"\n")
    }
    out.close
  } 
}

